from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from typing import Dict, Any, List, Optional
import uvicorn
import os
from dotenv import load_dotenv
import time
import json

# 加载环境变量
load_dotenv()

app = FastAPI(
    title="LangGraph FastAPI Demo",
    description="展示LangGraph各种功能的完整演示项目",
    version="1.0.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 挂载静态文件
if os.path.exists("static"):
    app.mount("/static", StaticFiles(directory="static"), name="static")

# 请求模型
class GraphRequest(BaseModel):
    input_data: Dict[str, Any]
    config: Optional[Dict[str, Any]] = None

class GraphResponse(BaseModel):
    result: Dict[str, Any]
    execution_time: float
    graph_type: str

# 简化的演示类
class BasicGraphDemo:
    """基础图执行演示"""
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行基础图"""
        message = input_data.get("message", "你好")
        
        # 模拟图执行流程
        steps = [
            f"1. 处理输入: {message}",
            f"2. 生成响应: AI助手回复: 我理解您说的是 '{message}'，我会为您提供帮助。",
            f"3. 格式化输出: 最终结果: AI助手回复: 我理解您说的是 '{message}'，我会为您提供帮助。"
        ]
        
        return {
            "final_state": {"messages": [message], "result": steps[-1]},
            "messages": [message, steps[1]],
            "final_result": steps[-1],
            "execution_steps": steps
        }

class ConditionalGraphDemo:
    """条件分支图演示"""
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行条件分支图"""
        message = input_data.get("message", "请帮我解答一个问题")
        
        # 简单的决策逻辑
        if "问题" in message or "解答" in message or "?" in message:
            decision = "question"
            path = "analyzed -> decided: question -> question_handled"
            result = "这是问题处理路径的响应"
        elif "任务" in message or "执行" in message or "做" in message:
            decision = "task"
            path = "analyzed -> decided: task -> task_handled"
            result = "这是任务处理路径的响应"
        else:
            decision = "chat"
            path = "analyzed -> decided: chat -> chat_handled"
            result = "这是聊天处理路径的响应"
        
        return {
            "final_state": {"decision": decision, "path": path},
            "decision": decision,
            "execution_path": path,
            "final_result": f"最终结果: {result} (路径: {path})"
        }

class ParallelGraphDemo:
    """并行执行图演示"""
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行并行执行图"""
        import time
        
        # 模拟并行任务
        tasks = {
            "task_a": "数据处理完成",
            "task_b": "网络请求完成", 
            "task_c": "文件操作完成"
        }
        
        # 模拟并行执行时间
        time.sleep(0.2)
        
        return {
            "final_state": {"parallel_results": tasks},
            "parallel_results": tasks,
            "combined_result": f"合并结果: {tasks}",
            "final_result": f"并行执行完成: 合并结果: {tasks}"
        }

class StateManagementDemo:
    """状态管理演示"""
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行状态管理图"""
        message = input_data.get("message", "开始新会话")
        
        # 模拟状态管理
        session_id = f"session_{hash(message)}"
        user_preferences = {}
        
        if "喜欢" in message:
            user_preferences["likes"] = message
        elif "不喜欢" in message:
            user_preferences["dislikes"] = message
        
        session_summary = {
            "session_id": session_id,
            "total_steps": 6,
            "user_preferences": user_preferences,
            "execution_history": ["会话初始化", f"处理输入: {message}", "更新用户信息", "跟踪活动: 步骤3", "生成响应", "保存会话状态"]
        }
        
        return {
            "final_state": {"session_summary": session_summary},
            "session_summary": session_summary,
            "final_result": f"会话已保存: {session_summary}"
        }

class ToolCallingDemo:
    """工具调用演示"""
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行工具调用图"""
        import datetime
        
        # 模拟工具调用结果
        tool_results = {
            "weather": {
                "location": "北京",
                "temperature": "25°C",
                "condition": "晴天",
                "humidity": "60%",
                "timestamp": datetime.datetime.now().isoformat()
            },
            "calculator": {
                "expression": "2 + 3 * 4",
                "result": 14,
                "steps": ["2 + 3 * 4", "2 + 12", "14"],
                "timestamp": datetime.datetime.now().isoformat()
            },
            "translator": {
                "source_text": "Hello, how are you?",
                "target_language": "中文",
                "translated_text": "你好，你好吗？",
                "confidence": 0.95,
                "timestamp": datetime.datetime.now().isoformat()
            }
        }
        
        final_response = """=== 工具调用演示结果 ===
🌤️ 天气信息: 北京 25°C 晴天
🧮 计算结果: 2 + 3 * 4 = 14
🌐 翻译结果: Hello, how are you? → 你好，你好吗？"""
        
        return {
            "final_state": {"tool_results": tool_results},
            "tool_results": tool_results,
            "tools_used": list(tool_results.keys()),
            "final_result": final_response
        }

class LoopGraphDemo:
    """循环图演示"""
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行循环图"""
        max_iterations = input_data.get("max_iterations", 5)
        
        # 模拟循环执行
        loop_history = ["循环初始化"]
        for i in range(1, max_iterations + 1):
            loop_history.append(f"迭代{i}: 第{i}次迭代处理完成")
            loop_history.append(f"检查条件: 迭代{i}, 继续={i < max_iterations}")
        
        loop_history.append("循环结束")
        
        return {
            "final_state": {"iteration_count": max_iterations, "loop_history": loop_history},
            "total_iterations": max_iterations,
            "max_iterations": max_iterations,
            "loop_history": loop_history,
            "final_result": f"循环完成，共执行{max_iterations}次迭代"
        }

class ErrorHandlingDemo:
    """错误处理演示"""
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行错误处理图"""
        import random
        
        max_retries = input_data.get("max_retries", 3)
        retry_count = 0
        error_history = ["初始化开始"]
        
        # 模拟重试逻辑
        for attempt in range(max_retries + 1):
            retry_count = attempt
            success_rate = 0.3 + (attempt * 0.2)
            
            if random.random() < success_rate:
                error_history.append(f"操作成功，重试{attempt}次")
                break
            else:
                error_history.append(f"操作失败: 操作失败，这是第{attempt + 1}次尝试")
                if attempt < max_retries:
                    error_history.append(f"准备第{attempt + 1}次重试")
                else:
                    error_history.append("执行备用操作")
        
        error_history.append("处理完成")
        
        if retry_count < max_retries:
            final_result = f"处理完成，操作成功 (重试{retry_count}次)"
        else:
            final_result = f"处理完成，经过{retry_count}次重试后，使用备用方案"
        
        return {
            "final_state": {"retry_count": retry_count, "error_history": error_history},
            "total_retries": retry_count,
            "max_retries": max_retries,
            "error_info": {} if retry_count < max_retries else {"error_type": "OperationFailed"},
            "error_history": error_history,
            "final_result": final_result
        }

# 初始化演示实例
basic_demo = BasicGraphDemo()
conditional_demo = ConditionalGraphDemo()
parallel_demo = ParallelGraphDemo()
state_demo = StateManagementDemo()
tool_demo = ToolCallingDemo()
loop_demo = LoopGraphDemo()
error_demo = ErrorHandlingDemo()

@app.get("/")
async def root():
    """根路径，返回项目信息"""
    return {
        "message": "LangGraph FastAPI Demo",
        "version": "1.0.0",
        "available_endpoints": [
            "/basic-graph",
            "/conditional-graph", 
            "/parallel-graph",
            "/state-management",
            "/tool-calling",
            "/loop-graph",
            "/error-handling",
            "/docs",
            "/demo"
        ]
    }

@app.get("/demo")
async def demo_page():
    """演示页面"""
    if os.path.exists("static/index.html"):
        return FileResponse("static/index.html")
    else:
        return {"message": "Demo page not found. Please check if static/index.html exists."}

@app.post("/basic-graph", response_model=GraphResponse)
async def basic_graph_demo(request: GraphRequest):
    """基础图执行演示"""
    try:
        start_time = time.time()
        
        result = basic_demo.run(request.input_data, request.config)
        
        execution_time = time.time() - start_time
        
        return GraphResponse(
            result=result,
            execution_time=execution_time,
            graph_type="basic_graph"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/conditional-graph", response_model=GraphResponse)
async def conditional_graph_demo(request: GraphRequest):
    """条件分支图演示"""
    try:
        start_time = time.time()
        
        result = conditional_demo.run(request.input_data, request.config)
        
        execution_time = time.time() - start_time
        
        return GraphResponse(
            result=result,
            execution_time=execution_time,
            graph_type="conditional_graph"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/parallel-graph", response_model=GraphResponse)
async def parallel_graph_demo(request: GraphRequest):
    """并行执行图演示"""
    try:
        start_time = time.time()
        
        result = parallel_demo.run(request.input_data, request.config)
        
        execution_time = time.time() - start_time
        
        return GraphResponse(
            result=result,
            execution_time=execution_time,
            graph_type="parallel_graph"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/state-management", response_model=GraphResponse)
async def state_management_demo(request: GraphRequest):
    """状态管理演示"""
    try:
        start_time = time.time()
        
        result = state_demo.run(request.input_data, request.config)
        
        execution_time = time.time() - start_time
        
        return GraphResponse(
            result=result,
            execution_time=execution_time,
            graph_type="state_management"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/tool-calling", response_model=GraphResponse)
async def tool_calling_demo(request: GraphRequest):
    """工具调用演示"""
    try:
        start_time = time.time()
        
        result = tool_demo.run(request.input_data, request.config)
        
        execution_time = time.time() - start_time
        
        return GraphResponse(
            result=result,
            execution_time=execution_time,
            graph_type="tool_calling"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/loop-graph", response_model=GraphResponse)
async def loop_graph_demo(request: GraphRequest):
    """循环图演示"""
    try:
        start_time = time.time()
        
        result = loop_demo.run(request.input_data, request.config)
        
        execution_time = time.time() - start_time
        
        return GraphResponse(
            result=result,
            execution_time=execution_time,
            graph_type="loop_graph"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/error-handling", response_model=GraphResponse)
async def error_handling_demo(request: GraphRequest):
    """错误处理演示"""
    try:
        start_time = time.time()
        
        result = error_demo.run(request.input_data, request.config)
        
        execution_time = time.time() - start_time
        
        return GraphResponse(
            result=result,
            execution_time=execution_time,
            graph_type="error_handling"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """健康检查端点"""
    return {"status": "healthy", "message": "LangGraph FastAPI Demo is running"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
